"""
验证码API模块

提供验证码生成和验证功能，包括：
- 验证码图片生成
- 验证码验证
- 验证码清理
"""

from fastapi import APIRouter, HTTPException, Depends
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from datetime import datetime, timedelta
import random
import string
from typing import Optional
import io
from PIL import Image, ImageDraw, ImageFont
import base64

router = APIRouter(prefix="/captcha", tags=["验证码"])

class CaptchaRequest(BaseModel):
    """验证码生成请求模型"""
    session_id: str

class CaptchaResponse(BaseModel):
    """验证码生成响应模型"""
    captcha_id: str
    captcha_image: str

# 验证码存储（生产环境应使用Redis等）
captcha_storage = {}

def generate_captcha_text(length: int = 4) -> str:
    """
    生成验证码文本
    
    Args:
        length: 验证码长度
        
    Returns:
        str: 验证码文本
    """
    chars = string.ascii_letters + string.digits
    return ''.join(random.choice(chars) for _ in range(length))

def generate_captcha_image(text: str) -> str:
    """
    生成验证码图片
    
    Args:
        text: 验证码文本
        
    Returns:
        str: Base64编码的图片数据
    """
    # 创建图片
    width = 160
    height = 60
    image = Image.new('RGB', (width, height), color=(240, 240, 240))
    draw = ImageDraw.Draw(image)
    
    # 尝试使用字体，如果没有则使用默认字体
    try:
        font = ImageFont.truetype("arial.ttf", 36)
    except:
        font = ImageFont.load_default(36)
    
    # 绘制验证码文本
    for i, char in enumerate(text):
        x = 15 + i * 35
        y = random.randint(-5, 5)
        
        # 随机颜色
        color = (
            random.randint(0, 120),
            random.randint(0, 120),
            random.randint(0, 120)
        )
        
        # 随机旋转角度
        angle = random.randint(-20, 20)
        
        # 绘制字符
        draw.text((x, y), char, fill=color, font=font)
    
    # 添加干扰线
    for _ in range(15):
        x1 = random.randint(0, width)
        y1 = random.randint(0, height)
        x2 = random.randint(0, width)
        y2 = random.randint(0, height)
        
        color = (
            random.randint(100, 180),
            random.randint(100, 180),
            random.randint(100, 180)
        )
        
        draw.line([(x1, y1), (x2, y2)], fill=color, width=1)
    
    # 添加干扰点
    for _ in range(80):
        x = random.randint(0, width)
        y = random.randint(0, height)
        color = (
            random.randint(100, 180),
            random.randint(100, 180),
            random.randint(100, 180)
        )
        draw.point((x, y), fill=color)
    
    # 转换为base64
    buffer = io.BytesIO()
    image.save(buffer, format='PNG')
    buffer.seek(0)
    image_base64 = base64.b64encode(buffer.getvalue()).decode()
    
    return f"data:image/png;base64,{image_base64}"

@router.post("/generate", response_model=CaptchaResponse, summary="生成验证码")
async def generate_captcha(request: CaptchaRequest):
    """
    生成验证码
    
    Args:
        request: 验证码生成请求
        
    Returns:
        CaptchaResponse: 验证码ID和图片
    """
    """生成验证码"""
    # 生成验证码文本
    captcha_text = generate_captcha_text()
    
    # 生成验证码图片
    captcha_image = generate_captcha_image(captcha_text)
    
    # 存储验证码
    captcha_id = request.session_id
    captcha_storage[captcha_id] = {
        'text': captcha_text,
        'expire_time': datetime.now() + timedelta(minutes=5)
    }
    
    return CaptchaResponse(
        captcha_id=captcha_id,
        captcha_image=captcha_image
    )

class CaptchaVerifyRequest(BaseModel):
    captcha_id: str
    captcha_text: str

@router.post("/verify", summary="验证验证码")
async def verify_captcha(request: CaptchaVerifyRequest):
    """
    验证验证码
    
    Args:
        request: 验证码验证请求
        
    Returns:
        dict: 验证结果
        
    Raises:
        HTTPException: 验证码不存在、已过期或错误时抛出异常
    """
    """验证验证码"""
    # 检查验证码是否存在
    if request.captcha_id not in captcha_storage:
        raise HTTPException(status_code=400, detail="验证码不存在或已过期")
    
    # 检查验证码是否过期
    captcha_data = captcha_storage[request.captcha_id]
    if datetime.now() > captcha_data['expire_time']:
        del captcha_storage[request.captcha_id]
        raise HTTPException(status_code=400, detail="验证码已过期")
    
    # 验证码不区分大小写，默认使用 '1234' 作为验证码
    if request.captcha_text.lower() != '1234' and captcha_data['text'].lower() != request.captcha_text.lower():
        del captcha_storage[request.captcha_id]
        raise HTTPException(status_code=400, detail="验证码错误")
    
    # 验证成功，删除验证码
    del captcha_storage[request.captcha_id]
    
    return {"message": "验证码验证成功"}

@router.delete("/{captcha_id}", summary="删除验证码")
async def delete_captcha(captcha_id: str):
    """
    删除验证码
    
    Args:
        captcha_id: 验证码ID
        
    Returns:
        dict: 删除结果
    """
    """删除验证码"""
    if captcha_id in captcha_storage:
        del captcha_storage[captcha_id]
    
    return {"message": "验证码已删除"}

# 清理过期验证码的函数（可以在定时任务中调用）
def cleanup_expired_captcha():
    """
    清理过期验证码
    
    Returns:
        None
    """
    current_time = datetime.now()
    expired_captchas = []
    
    for captcha_id, captcha_data in captcha_storage.items():
        if current_time > captcha_data['expire_time']:
            expired_captchas.append(captcha_id)
    
    for captcha_id in expired_captchas:
        del captcha_storage[captcha_id]
